from sympy import Symbol, Eq, Abs

from modulus.solver import Solver
from modulus.dataset import TrainDomain, ValidationDomain
from modulus.data import Validation
from modulus.sympy_utils.functions import parabola
from modulus.sympy_utils.geometry_2d import Rectangle, Line
from modulus.csv_utils.csv_rw import csv_to_dict
from modulus.PDES.navier_stokes import NavierStokes, IntegralContinuity
from modulus.controller import ModulusController

# 域参数
height = 0.1
width = 0.1
inlet_vel = 1.0

# 定义几何图形
rec = Rectangle((-width / 2, -height / 2), (width / 2, height / 2))
obstacle = Line((0, -height/4), (0, height/4), 1)
obstacle.rotate(1.05)
geo = rec

plane1 = Line((-3*width/8, -height/2),
(-3*width/8, height/2), 1)
plane2 = Line((-2*width/8, -height/2),
(-2*width/8, height/2), 1)
plane3 = Line((-width/8, -height/2),
(-width/8, height/2), 1)
plane4 = Line((0, -height/2),
(0, height/2), 1)
plane5 = Line((width/8, -height/2),
(width/8, height/2), 1)
plane6 = Line((2*width/8, -height/2),
(2*width/8, height/2), 1)
plane7 = Line((3*width/8, -height/2),
(3*width/8, height/2), 1)

# 定义sympy变量以参数化域曲线
x, y = Symbol('x'), Symbol('y')

class LDCTrain(TrainDomain):
  def __init__(self, **config):
    super(LDCTrain, self).__init__()

    parabola_sympy = parabola(y, inter_1=-height/2, inter_2=height/2, height=inlet_vel)
    inletBC = geo.boundary_bc(outvar_sympy={'u': parabola_sympy, 'v': 0},
                              batch_size_per_area=64,
                              criteria=Eq(x, -width/2))
    self.add(inletBC, name="Inlet")

    outletBC = geo.boundary_bc(outvar_sympy={'p': 0, 'integral_continuity':0.1333333},
                               batch_size_per_area=64,
                               criteria=Eq(x, width/2))
    self.add(outletBC, name="Outlet")

    topWall = geo.boundary_bc(outvar_sympy={'u': 0, 'v': 0},
                              batch_size_per_area=256,
                              criteria=Eq(y, height/2))
    self.add(topWall, name="TopWall")

    bottomWall = geo.boundary_bc(outvar_sympy={'u': 0, 'v': 0},
                                 batch_size_per_area=256,
                                 criteria=Eq(y, -height/2))
    self.add(bottomWall, name="BottomWall")

    obstacleLine = obstacle.boundary_bc(outvar_sympy = {'u': 0, 'v': 0},
                                   batch_size_per_area=1000,
                                   lambda_sympy={'lambda_u':10, 'lambda_v':10})
    self.add(obstacleLine, name="obstacleLine")

    interior = geo.interior_bc(outvar_sympy={'continuity': 0, 'momentum_x': 0, 'momentum_y': 0},
                               bounds={x: (-width / 2, width / 2),
                                       y: (-height / 2, height / 2)},
                               lambda_sympy={'lambda_continuity': 10,
                                             'lambda_momentum_x': geo.sdf,
                                             'lambda_momentum_y': geo.sdf},
                               batch_size_per_area=10000)
    self.add(interior, name="Interior")

    plane1Cont = plane1.boundary_bc(outvar_sympy={'integral_continuity':0.1333333},
                                    batch_size_per_area=256,
                                    lambda_sympy={'lambda_integral_continuity':10})
    plane2Cont = plane2.boundary_bc(outvar_sympy={'integral_continuity':0.1333333},
                                    batch_size_per_area=256,
                                    lambda_sympy={'lambda_integral_continuity':10})
    plane3Cont = plane3.boundary_bc(outvar_sympy={'integral_continuity':0.1333333},
                                    batch_size_per_area=256,
                                    lambda_sympy={'lambda_integral_continuity':10})
    plane4Cont = plane4.boundary_bc(outvar_sympy={'integral_continuity':0.1333333},
                                    batch_size_per_area=256,
                                    lambda_sympy={'lambda_integral_continuity':10})
    plane5Cont = plane5.boundary_bc(outvar_sympy={'integral_continuity':0.1333333},
                                    batch_size_per_area=256,
                                    lambda_sympy={'lambda_integral_continuity':10})
    plane6Cont = plane6.boundary_bc(outvar_sympy={'integral_continuity':0.1333333},
                                    batch_size_per_area=256,
                                    lambda_sympy={'lambda_integral_continuity':10})
    plane7Cont = plane7.boundary_bc(outvar_sympy={'integral_continuity':0.1333333},
                                    batch_size_per_area=256,
                                    lambda_sympy={'lambda_integral_continuity':10})


    self.add(plane1Cont, name="integralContinuity1")
    self.add(plane2Cont, name="integralContinuity2")
    self.add(plane3Cont, name="integralContinuity3")
    self.add(plane4Cont, name="integralContinuity4")
    self.add(plane5Cont, name="integralContinuity5")
    self.add(plane6Cont, name="integralContinuity6")
    self.add(plane7Cont, name="integralContinuity7")

# 校验数据
mapping = {'Points:0': 'x', 'Points:1': 'y', 'U:0': 'u', 'U:1': 'v', 'p': 'p'}
openfoam_var = csv_to_dict('openfoam/cavity_uniformVel0.csv', mapping)
openfoam_var['x'] += -width / 2  # 中心OpenFoam数据
openfoam_var['y'] += -height / 2  # 中心OpenFoam数据

openfoam_invar_numpy = {key: value for key, value in openfoam_var.items() if key in ['x', 'y']}
openfoam_outvar_numpy = {key: value for key, value in openfoam_var.items() if key in ['u', 'v']}

class LDCVal(ValidationDomain):
  def __init__(self, **config):
    super(LDCVal, self).__init__()
    val = Validation.from_numpy(openfoam_invar_numpy, openfoam_outvar_numpy)
    self.add(val, name='Val')

class LDCSolver(Solver):
  train_domain = LDCTrain
  val_domain = LDCVal

  def __init__(self, **config):
    super(LDCSolver, self).__init__(**config)
    self.equations = (NavierStokes(nu=0.01, rho=1.0, dim=2,time=False).make_node()
                      +IntegralContinuity().make_node())
    flow_net = self.arch.make_node(name='flow_net',
                                   inputs=['x', 'y'],
                                   outputs=['u', 'v', 'p'])
    self.nets = [flow_net]

  @classmethod
  def update_defaults(cls, defaults):
    defaults.update({
        'network_dir': './network_checkpoint_ldc_2d',
        'decay_steps': 4000,
        'max_steps': 400000
    })


if __name__ == '__main__':
  ctr = ModulusController(LDCSolver)
  ctr.run()
